from kmeans import *
from keras.datasets import mnist

# 导入mnist数据集
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

batch = 10000
clusters = 20
train_matrix = np.zeros((batch, 784))
train_labels = train_labels[0:batch]

for i in range(batch):
    for j in range(28):
        for k in range(28):
            train_matrix[i, 28 * j + k] = train_images[i][j][k]

a = kmeans(clusters, 300, train_matrix, train_labels)
a.fit()

label_num = np.zeros((clusters, 10))

for i in range(a.classifications.shape[0]):
    pred = int(a.classifications[i])
    truth = int(train_labels[i])
    label_num[pred][truth] += 1

label2num = label_num.argmax(axis=1)
set(label2num)
train_preds = np.zeros(train_labels.shape)
for i in range(train_preds.shape[0]):
    train_preds[i] = label2num[a.classifications[i]]

print("mnist数据集10000个数据聚类准确率", (np.sum(train_preds == train_labels) / train_labels.shape[0]))

